"""
test trained model on given dataset, save all results into text files for mAP calculation.
"""
import time
import os
import torch
from .inference import ModelInference


class Tester(object):
    def __init__(self, test_opts):
        self.model_inference = ModelInference(test_opts)
    
    def test(self, image_folder, text_folder):
        """
        test on a folder of images and save results into separate folder
        """
        
        if not os.path.exists(text_folder):
            os.makedirs(text_folder)
        
        for image_name in os.listdir(image_folder):
            image_path = os.path.join(image_folder, image_name)
            boxes, labels = self.model_inference.detect(image_path)

            # save into file
            base_name = '.'.join(image_name.split('.')[:-1])
            text_path = os.path.join(text_folder, '{}.txt'.format(base_name))
            with open(text_path) as f:
                for i, box in enumerate(boxes):
                    box = box.aslist()
                    box = [int(v) for v in box]
                    label = labels[i]

                    f.write('{} {} {} {} {}\n'.format(label, box[0], box[1], box[2], box[3]))
